import cv2
import numpy as np

'''滤波'''
# 高斯滤波
def gauss(img, ksize: int, sigmaX: int, sigmaY: int):
    gauss_img = cv2.GaussianBlur(img, (ksize, ksize), sigmaX, sigmaY)
    return gauss_img, {'ksize': ksize, 'sigmaX': sigmaX, 'sigmaY': sigmaY}

# 中值滤波
def median(img, k: int):
    median_img = cv2.medianBlur(img, k)
    return median_img, {'k': k}

# 均值滤波
def mean(img, ksize: int):
    mean_img = cv2.blur(img, (ksize, ksize))
    return mean_img, {'ksize': ksize}

# 双边滤波
def bilateral(img, d: int, sigmaColor: int, sigmaSpace: int):
    bilatera_img = cv2.bilateralFilter(img, d, sigmaColor, sigmaSpace)
    return bilatera_img, {'d': d, 'sigmaColor': sigmaColor, 'sigmaSpace': sigmaSpace}

# 方框滤波
def box(img, depth: int, ksize: int, normalize: bool):
    box_img = cv2.boxFilter(img, depth, (ksize, ksize), normalize)
    return box_img, {'depth': depth, 'ksize': ksize, 'normalize': normalize}

# 2D卷积
def filter2D(img, depth: int, kernel: np.ndarray):
    filter2D_img = cv2.filter2D(img, depth, kernel)
    return filter2D_img, {'depth': depth, 'kernel': kernel}

def blur(img, **params):
    blur_type = params['blur_type']
    print(f"blur type is: {blur_type}")
    params.pop('blur_type')
    # 将params中的字符串数字转换为整数
    for key in params:
        if params[key].isdigit():
            params[key] = int(params[key])
    return FILTER[blur_type](img, **params)

'''去噪'''
# 非局部平均去噪（灰度）
def fastNlMeans(img, h: int, templateWindowSize: int, searchWindowSize: int):
    nlm_img = cv2.fastNlMeansDenoising(img, h=h, templateWindowSize=templateWindowSize, searchWindowSize=searchWindowSize)
    return nlm_img, {'h': h,
                     'templateWindowSize': templateWindowSize, 'searchWindowSize': searchWindowSize}

# 非局部平均去噪（彩图）
def fastNlMeansColored(img, h: int, hColor: int, templateWindowSize: int, searchWindowSize: int):
    nlmc_img = cv2.fastNlMeansDenoisingColored(img, h=h, hColor=hColor, templateWindowSize=templateWindowSize, searchWindowSize=searchWindowSize)
    return nlmc_img, {'h': h, 'hForColorComponents': hColor,
                      'templateWindowSize': templateWindowSize, 'searchWindowSize': searchWindowSize}

# 非局部平均去噪（灰度帧序列）
def fastNlMeansMulti(img, index: int, temporalWindowSize: int,
                     h: int, templateWindowSize: int, searchWindowSize: int):
    nlmm_img = cv2.fastNlMeansDenoisingMulti(img, imgToDenoiseIndex=index, temporalWindowSize=temporalWindowSize,
                                             h=h, templateWindowSize=templateWindowSize, searchWindowSize=searchWindowSize)
    return nlmm_img, {'index': index, 'temporalWindowSize': temporalWindowSize,
                      'h': h, 'templateWindowSize': templateWindowSize, 'searchWindowSize': searchWindowSize}

# 非局部平均去噪（彩色帧序列）
def fastNlMeansColoredMulti(img, index: int, temporalWindowSize: int,
                            h: int, templateWindowSize: int, searchWindowSize: int):
    nlmm_img = cv2.fastNlMeansDenoisingColoredMulti(img, imgToDenoiseIndex=index, temporalWindowSize=temporalWindowSize,
                                                    h=h, templateWindowSize=templateWindowSize, searchWindowSize=searchWindowSize)
    return nlmm_img, {'index': index, 'temporalWindowSize': temporalWindowSize,
                      'h': h, 'templateWindowSize': templateWindowSize, 'searchWindowSize': searchWindowSize}

def denosing(img, **params):
    denosing_type = params['denosing_type']
    print(f"denosing type is: {denosing_type}")
    params.pop('denosing_type')
    # 将params中的字符串数字转换为整数
    for key in params:
        if params[key].isdigit():
            params[key] = int(params[key])
    return DENOISE[denosing_type](img, **params)


'''形态学'''
OP = {
    'OPEN': cv2.MORPH_OPEN,
    'CLOSE': cv2.MORPH_CLOSE,
    'GRADIENT': cv2.MORPH_GRADIENT,
    'TOPHAT': cv2.MORPH_TOPHAT,
    'BLACKHAT': cv2.MORPH_BLACKHAT
}

# 腐蚀
def erode(img, kernel: np.ndarray, iterations: int = 1):
    erosion_img = cv2.erode(img, kernel, iterations)
    return erosion_img, {'iterations': iterations}

# 膨胀
def dilate(img, kernel: np.ndarray, iterations: int = 1):
    dilation_img = cv2.dilate(img, kernel, iterations)
    return dilation_img, {'iterations': iterations}

# 开闭运算、梯度、礼帽、黑帽
def morphologyEx(img, op, kernel: np.ndarray):
    mor_img = cv2.morphologyEx(img, OP[op], kernel)
    return mor_img, {'op': op}

def morphology(img, **params):
    morphology_type = params['morphology_type']
    print(f"morphology type is: {morphology_type}")
    params.pop('morphology_type')
    # 将params中的字符串数字转换为整数
    for key in params:
        if params[key].isdigit():
            params[key] = int(params[key])
    params['kernel'] = cv2.getStructuringElement(cv2.MORPH_RECT, (params['kernel'], params['kernel']))
    return MORPHOLOGY[morphology_type](img, **params)

# 滤波
FILTER = {
    'gauss': gauss,
    'median': median,
    'mean': mean,
    'bilateral': bilateral,
    'box': box,
    'filter2D': filter2D
}

# 去噪
DENOISE = {
    'fastNlMeans': fastNlMeans,
    'fastNlMeansColored': fastNlMeansColored,
    'fastNlMeansMulti': fastNlMeansMulti,
    'fastNlMeansColoredMulti': fastNlMeansColoredMulti
}

# 形态学
MORPHOLOGY = {
    'erode': erode,
    'dilate': dilate,
    'morphologyEx': morphologyEx,
}

#旋转
def rotate(img, rotateCode: int) :

    if rotateCode==90:#顺时针90度
        img=cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    elif rotateCode==180:#顺时针180度
        img=cv2.rotate(img, cv2.ROTATE_180)
    elif rotateCode==270:##顺时针270度（逆时针90度）
        img=cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    return img, {}
#裁剪
def clip(img, x1:int,x2: int,y1: int,y2: int):
    x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)
    img=img[x1:x2,y1:y2]
    return img, {}

THRESHOLDTYPE = {
        'binary': cv2.THRESH_BINARY,
        'binary-inv': cv2.THRESH_BINARY_INV,
        'trunc': cv2.THRESH_TRUNC,
        'tozero': cv2.THRESH_TOZERO,
        'tozero-inv': cv2.THRESH_TOZERO_INV,
        'mask': cv2.THRESH_MASK,
        'otsu': cv2.THRESH_OTSU,
        'triangle': cv2.THRESH_TRIANGLE
        }

def threshold(img,thresh:int,maxval:int,thresholdtype): #阈值化
    retval, img=cv2.threshold(img,int(thresh), int(maxval), THRESHOLDTYPE[thresholdtype])
    return img,{'thresh':thresh,'maxval':maxval,'thresholdtype':thresholdtype}

def resize(img,dsize:int,fx:int,fy:int,interpolation:int=cv2.INTER_LINEAR):
    img=cv2.resize(img,dsize,fx,fy,interpolation)
    return img,{'dsize':dsize,'fx':fx,'fy':fy,'interpolation':interpolation}

def cvtColor(img,code:int,dstCn:int=0):
    img=cv2.cvtColor(img,img,code,dstCn)
    return img,{'code':code,'dstCn':dstCn}
